from core.function import Function
from core.cuda import cuda

# 均方误差
# 先求出N个点的每个点的平方误差，再将它们加起来
class MeanSquaredError(Function):
    def forward(self, x0, x1):
        """
        均方误差的前向传播。

        Args:
            x0 (Variable): 输入变量。
            x1 (Variable): 输入变量。

        Returns:
            Variable: 输出变量。
        """
        x0, x1 = cuda.to_array(x0), cuda.to_array(x1)
        diff = x0 - x1
        y = (diff ** 2).sum() / len(diff)
        return y

    def backward(self, gy):
        """
        均方误差的反向传播。

        Args:
            gy (Variable): 相对于输出的梯度。

        Returns:
            tuple: 包含相对于输入 x0 和 x1 的梯度的元组。
        """
        x0, x1 = self.inputs
        diff = x0 - x1
        gx0 = gy * diff * (2. / len(diff))
        gx1 = -gx0
        return gx0, gx1

    def mean_squared_error(self, x0, x1):
        """
        均方误差操作。

        Args:
            x0 (Variable): 输入变量。
            x1 (Variable): 输入变量。

        Returns:
            Variable: 输出变量。
        """
        return self(x0, x1)

